#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/4/3
# @Author  : geekhch
# @Email   : geekhch@qq.com
# @Desc    : None

import sys, json
import torch
import os
import numpy as np
from OpenNRE import opennre
from OpenNRE.opennre import encoder, model, framework
from myPath import *
from myVocab import vocab
from utils import logger

# Some basic settings
if not os.path.exists(DULE_MODEL):
    os.makedirs(DULE_MODEL)

from datetime import datetime
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
ckpt = path.join(DULE_MODEL, f'{current_time}_cnn.pth.tar')

rel2id = json.load(open(os.path.join(DULE_DIR, 'rel2id.json')))
wordi2d = vocab.get_word2id_dict()
word2vec = vocab.get_embedding()

# Define the sentence encoder
sentence_encoder = opennre.encoder.CNNEncoder(
    token2id=wordi2d,
    max_length=60,
    word_size=word2vec.shape[1],
    position_size=5,
    hidden_size=800,
    blank_padding=True,
    kernel_size=3,
    padding_size=1,
    word2vec=word2vec.astype(np.float),
    dropout=0.5
)

# Define the model
model = opennre.model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
logger.info(str(model))

# Define the whole training framework
def get_framework():
    framework = opennre.framework.SentenceRE(
        train_path=os.path.join(DULE_DIR, 'train.txt'),
        val_path=os.path.join(DULE_DIR, 'dev.txt'),
        test_path=os.path.join(DULE_DIR, 'dev.txt'),
        model=model,
        ckpt=ckpt,
        batch_size=64,
        max_epoch=100,
        lr=0.1,
        weight_decay=1e-5,
        opt='sgd'
    )
    return framework

if __name__ == '__main__':
    # Train the model
    # framework.train_model()

    # Test the model
    framework = get_framework()
    framework.load_state_dict(torch.load(f'{DULE_MODEL}/Apr11_13-29-53_cnn.pth.tar')['state_dict'])
    result = framework.eval_model(framework.test_loader)
    print(result)

    # model.load_state_dict(torch.load(f'{DULE_MODEL}/Apr11_13-29-53.pth.tar', map_location='cpu')['state_dict'])
    # result = model.infer({"token": ["《", "步", "步", "惊", "心", "》", "改", "编", "自", "著", "名", "作", "家", "桐", "华", "的", "同", "名", "清", "穿", "小", "说", "《", "甄", "嬛", "传", "》", "改", "编", "自", "流", "潋", "紫", "所", "著", "的", "同", "名", "小", "说", "电", "视", "剧", "《", "何", "以", "笙", "箫", "默", "》", "改", "编", "自", "顾", "漫", "同", "名", "小", "说", "《", "花", "千", "骨", "》", "改", "编", "自", "f", "r", "e", "s", "h", "果", "果", "同", "名", "小", "说", "《", "裸", "婚", "时", "代", "》", "是", "月", "影", "兰", "析", "创", "作", "的", "一", "部", "情", "感", "小", "说", "《", "琅", "琊", "榜", "》", "是", "根", "据", "海", "宴", "同", "名", "网", "络", "小", "说", "改", "编", "电", "视", "剧", "《", "宫", "锁", "心", "玉", "》", "，", "又", "名", "《", "宫", "》", "《", "雪", "豹", "》", "，", "该", "剧", "改", "编", "自", "网", "络", "小", "说", "《", "特", "战", "先", "驱", "》", "《", "我", "是", "特", "种", "兵", "》", "由", "红", "遍", "网", "络", "的", "小", "说", "《", "最", "后", "一", "颗", "子", "弹", "留", "给", "我", "》", "改", "编", "电", "视", "剧", "《", "来", "不", "及", "说", "我", "爱", "你", "》", "改", "编", "自", "匪", "我", "思", "存", "同", "名", "小", "说", "《", "来", "不", "及", "说", "我", "爱", "你", "》"], "h": {"name": "顾漫", "pos": [53, 55], "type": "人物"}, "t": {"name": "何以笙箫默", "pos": [44, 49], "type": "图书作品"}, "relation": "创作者"})
    # print(model)
